import torch
import torch.nn as nn


class yolo_v1_loss(nn.Module):
    """
    yolov1损失函数
    1)  负责检测物体的bounding_box的中心点定位误差  λ1 * 累加(是否负责检测 * ((x-x`)^2 + (y-y`)^2))  其中 x`为标注值
    2)  负责检测物体的bounding_box的宽高定位误差 λ2 * 累加(是否负责检测 * ((根号(w)-根号(w`))^2 + (根号(h)-根号(h`))^2))  其中 w`为标注值
    3)  负责检测物体的bounding_box的置信度误差  累加(是否负责检测 * (预测目标置信度 - 预测框iou真实框)^2)
    4)  不负责检测物体的bounding_box的置信度误差  λ3 * 累加(是否不负责检测 * (预测目标置信度)^2)
    5)  负责检测物体的bounding_box的分类误差  累加(是否负责检测 * 类别累加(p(c) - p(c`) ^ 2))

    其中λ1,λ2一般取5 λ3一般取0.5
    """

    def __init__(self, device, bounding_box_size=2, classes_num=20, lamda1=5, lamda2=5, lamda3=0.5) -> None:
        """
        Args:
            bounding_box_size: 同一个输出有几个预设bounding_box_size
            classes_num:
        """
        super().__init__()
        self.classes_num = classes_num
        self.bounding_box_size = bounding_box_size
        self.lamda1 = lamda1
        self.lamda2 = lamda2
        self.lamda3 = lamda3
        self.device = device

    def forward(self, pred: torch.tensor, target: torch.tensor):
        """
          output: [batch_size, 7 * 7 * (5 + 5 + 20)]
          ground_truth: [batch_size, 7 * 7 * 9]
          ground_mask_positive: [batch_size, 7, 7, 1]
        """
        ground_truth = target[0]
        ground_mask_positive = target[1]
        ground_mast_negative = torch.logical_not(target[1])
        batch_size = len(pred)
        mseLoss = nn.MSELoss(size_average=False)
        # 见当前loss类注释
        loss = torch.zeros(5)
        # 这个批次中所有 预测值中存在预测框的格子
        predict_positive_boxes = torch.masked_select(pred, ground_mask_positive).view(-1, self.bounding_box_size * 5 + self.classes_num)
        # 这个批次中所有图片 预测值中不存在预测框的格子 和上面的存在预测框的数据共 batch_size * 7 * 7 个预测框
        predict_negative_boxes = torch.masked_select(pred, ground_mast_negative).view(-1, self.bounding_box_size * 5 + self.classes_num)
        # 标注数据中存在预测框的网格
        positive_ground_truth = torch.masked_select(ground_truth, ground_mask_positive).view(-1, 5 + self.classes_num + 2)

        # 第一个预测框和第二预测框 并计算对应的iou
        one_predict_boxes = predict_positive_boxes[:, 0:5]
        two_predict_boxes = predict_positive_boxes[:, 5:10]
        boxes_one_iou = self.iou(one_predict_boxes, positive_ground_truth)
        boxes_two_iou = self.iou(two_predict_boxes, positive_ground_truth)

        # 哪个框的iou大,就用哪个框,哪个小,就不用哪个框
        positive_iou_bigest_box = torch.where(boxes_one_iou > boxes_two_iou, True, False)
        iou_box_len = len(positive_iou_bigest_box)
        positive_iou = torch.where(boxes_one_iou > boxes_two_iou, boxes_one_iou, boxes_two_iou)
        positive_box_mask = torch.zeros(iou_box_len, 10).to(self.device)
        nagetive_box_mask = torch.ones(iou_box_len, 10).to(self.device)

        for positive_iou_bigest_box_index in range(iou_box_len):
            if positive_iou_bigest_box[positive_iou_bigest_box_index]:
                positive_box_mask[positive_iou_bigest_box_index, 0:5] = torch.ones(5, )
                nagetive_box_mask[positive_iou_bigest_box_index, 0:5] = torch.zeros(5, )
            else:
                positive_box_mask[positive_iou_bigest_box_index, 5:10] = torch.ones(5, )
                nagetive_box_mask[positive_iou_bigest_box_index, 5:10] = torch.zeros(5, )

        positive_box_bigest_iou_box = torch.masked_select(predict_positive_boxes[:, 0:10], positive_box_mask.bool()).view(-1, 5)
        nagetive_box_bigest_iou_box = torch.masked_select(predict_positive_boxes[:, 0:10], nagetive_box_mask.bool()).view(-1, 5)

        # 置信度
        ground_class = positive_ground_truth[:, 5:  5 + self.classes_num]
        predict_class = predict_positive_boxes[:, self.bounding_box_size * 5: self.bounding_box_size * 5 + self.classes_num]

        #  ---------------------------------------- 公式 -----------------------------
        # 中心点定位误差
        center_point_error = self.lamda1 * torch.sum(torch.pow(positive_box_bigest_iou_box[:, 0] - positive_ground_truth[:, 0], 2) + torch.pow(positive_box_bigest_iou_box[:, 1] - positive_ground_truth[:, 1], 2))
        # 宽高定位误差
        width_and_height_positioning_error = self.lamda2 * torch.sum(torch.pow(torch.sqrt(positive_box_bigest_iou_box[:, 2] + 1e-8) - torch.sqrt(positive_ground_truth[:, 2] + 1e-8), 2) + torch.pow(torch.sqrt(positive_box_bigest_iou_box[:, 3] + 1e-8) - torch.sqrt(positive_ground_truth[:, 3] + 1e-8), 2))

        # 对的bbox的置信度误差
        confidence_error = torch.sum(torch.pow(positive_box_bigest_iou_box[:, 4] - positive_iou[:], 2))

        # 对不是bbox的置信度需要和0做误差
        nagetive_box = torch.cat([nagetive_box_bigest_iou_box, predict_negative_boxes[:, 0:5], predict_negative_boxes[:, 5:10]], dim=0)
        nagetive_box_confidence_error = self.lamda3 * torch.sum(torch.pow(nagetive_box[:, 4], 2))

        # 分类误差
        classification_error = mseLoss(predict_class, ground_class)
        loss[0] = center_point_error / batch_size
        loss[1] = width_and_height_positioning_error / batch_size
        loss[2] = confidence_error / batch_size
        loss[3] = nagetive_box_confidence_error / batch_size
        loss[4] = classification_error / batch_size
        loss_sum = torch.sum(loss)
        return loss_sum, (loss, predict_class, ground_class, positive_iou)

    def iou(self, box1, box2):
        """
        计算两个box的iou
        Args:
            box1: [batch, (center_x, center_y, ground_width, ground_height)]
            box2: [batch, (center_x, center_y, ground_width, ground_height)]
        Returns:
            iou
        """
        x1 = box1[:, 0] - box1[:, 2] / 2
        y1 = box1[:, 1] - box1[:, 3] / 2
        w1 = box1[:, 2]
        h1 = box1[:, 3]

        x2 = box2[:, 0] - box2[:, 2] / 2
        y2 = box2[:, 1] - box2[:, 3] / 2
        w2 = box2[:, 2]
        h2 = box2[:, 3]

        max_x1 = torch.max(x1, x2)
        max_y1 = torch.max(y1, y2)
        max_x2 = torch.min(x1 + w1, x2 + w2)
        max_y2 = torch.min(y1 + h1, y2 + h2)

        # 交集面积
        x_ = max_x2 - max_x1
        y_ = max_y2 - max_y1
        max_w = torch.where(x_ < 0, 0, x_)
        max_h = torch.where(y_ < 0, 0, y_)
        max_area = max_w * max_h

        # 并集面积
        # 计算并集面积
        union_area = w1 * h1 + w2 * h2 - max_area

        iou = max_area / union_area

        return iou
